import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from datasets import load_dataset

model_name = 'gpt2'  # 或者使用 'gpt2-medium', 'gpt2-large', 等
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

# 将模型设置为评估模式
model.eval()


def generate_response(prompt, max_length=100):
    # 编码输入文本
    inputs = tokenizer.encode(prompt, return_tensors='pt', skip_special_tokens=True, clean_up_tokenization_spaces=True)

    # 生成响应
    with torch.no_grad():  # 不需要计算梯度
        outputs = model.generate(inputs, max_length=max_length, num_return_sequences=1)

    # 解码生成的文本
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response


# 示例对话
user_input = "Hello! How are you?"
response = generate_response(user_input)
print("GPT-2:", response)
